/*
 * Copyright 2013 Evan Jones.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.keyczar;

import java.io.UnsupportedEncodingException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;

import javax.crypto.Cipher;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;

import org.keyczar.enums.CipherMode;
import org.keyczar.exceptions.Base64DecodingException;
import org.keyczar.exceptions.KeyczarException;
import org.keyczar.interfaces.EncryptedReader;
import org.keyczar.interfaces.KeyczarReader;
import org.keyczar.util.Base64Coder;
import org.keyczar.util.Util;

import com.google.gson.annotations.Expose;

/**
 * Reads metadata and keys, decrypted using a password. This is the Java
 * equivalent to C++'s KeysetPBEJSONFileReader:
 * https://code.google.com/p/keyczar/wiki/CppTutorial#Using_PBE_(password-based_encryption)_to_encrypt_keys
 */
public class KeyczarPBEReader implements EncryptedReader {
  private final KeyczarReader reader;
  private final String passphrase;

  // PBKDF2 NIST approved standard key derivation function
  private static final String PBKDF2_ALGORITHM = "PBKDF2WithHmacSHA1";

  // PBKDF2 RFC 2898 recommends at least 8 bytes (64 bits) of salt
  // http://tools.ietf.org/html/rfc2898#section-4
  // but NIST recommends at least 16 bytes (128 bits; see Section 5.1)
  // http://csrc.nist.gov/publications/nistpubs/800-132/nist-sp800-132.pdf
  static final int SALT_BYTES = 16;

  // NIST suggests count to be 1000 as a minimum, but that seems poor.
  // 4 GPUs can do 3M attempts/second with 1000 iterations. See:
  // http://blog.agilebits.com/2013/04/16/1password-hashcat-strong-master-passwords/
  // 10000 iterations; 7 mixed case random letters = 9 days to crack
  // 10000 iterations; 9 mixed case random letters = 193 years days to crack
  // C++ Keyczar uses 4096 iterations by default (crypto_factory.cc)
  @SuppressWarnings("unused")
  private static final int MIN_ITERATION_COUNT = 10000;

  // We use 50000 iterations by default to increase brute force difficulty
  // 7 mixed case random letters = 41 days to crack
  // 9 mixed case random letters = 867 years to crack
  static final int DEFAULT_ITERATION_COUNT = 50000;

  /** Generates a key using PBDKDF2, based on parameters in keySpec. */
  static byte[] pbkdf2(PBEKeySpec keySpec) {
    try {
      SecretKeyFactory factory = SecretKeyFactory.getInstance(PBKDF2_ALGORITHM);
      SecretKey key = factory.generateSecret(keySpec);
      return key.getEncoded();
    } catch (NoSuchAlgorithmException e) {
      throw new RuntimeException("Unexpected: unsupported key derivation function?", e);
    } catch (InvalidKeySpecException e) {
      throw new IllegalArgumentException("Invalid keySpec", e);
    }
  }

  private static final String AES_ALGORITHM = "AES";

  /**
   * Reads encrypted key files from the given reader and decrypts them
   * with the passphrase.
   *
   * @param reader The reader to read files from.
   * @param passphrase The passphrase used to derive the decryption key.
   */
  public KeyczarPBEReader(KeyczarReader reader, String passphrase) {
    this.reader = reader;
    this.passphrase = passphrase;
  }

  String decryptKey(String encryptedKeyData) throws KeyczarException {
    PBEKeyczarKey pbeMetadata = parsePBEMetadata(encryptedKeyData);

    try {
      Cipher decryptingCipher = initializePassphraseCipher(
          passphrase, pbeMetadata, Cipher.DECRYPT_MODE);

      byte[] encryptedKey = Base64Coder.decodeWebSafe(pbeMetadata.key);
      byte[] decrypted = decryptingCipher.doFinal(encryptedKey);
      return new String(decrypted, "UTF-8");
    } catch (java.security.GeneralSecurityException e) {
      throw new KeyczarException("Error decrypting PBE key", e);
    } catch (UnsupportedEncodingException e) {
      throw new RuntimeException("Should never occur");
    }
  }

  /**
   * Generate a Cipher from the passphrase, using parameters in pbeMetadata and iv.
   */
  private static Cipher initializePassphraseCipher(
      String passphrase, PBEKeyczarKey pbeMetadata, int cipherMode) throws
      NoSuchAlgorithmException, NoSuchPaddingException,
      InvalidKeyException, InvalidAlgorithmParameterException,
      Base64DecodingException {
    PBEKeySpec spec = keySpecFromJson(pbeMetadata, passphrase);
    byte[] keyBytes = pbkdf2(spec);
    SecretKeySpec key = new SecretKeySpec(keyBytes, AES_ALGORITHM);

    IvParameterSpec iv = new IvParameterSpec(Base64Coder.decodeWebSafe(pbeMetadata.iv));

    Cipher cipher = Cipher.getInstance(CipherMode.CBC.getMode());
    cipher.init(cipherMode, key, iv);
    return cipher;
  }

  @Override
  public String getKey() throws KeyczarException {
    return decryptKey(reader.getKey());
  }

  @Override
  public String getKey(int version) throws KeyczarException {
    return decryptKey((reader.getKey(version)));
  }

  @Override
  public String getMetadata() throws KeyczarException {
    return reader.getMetadata();
  }

  /** Gson struct for the key derivation parameters. */
  static final class PBEKeyczarKey {
    @Expose
    public String cipher;
    @Expose
    public String hmac;
    @Expose
    public int iterationCount;
    @Expose
    public String iv;
    @Expose
    public String key;
    @Expose
    public String salt;
  }

  /**
   * Encrypts key with password-based encryption.
   * @return a JSON string containing encrypted key and the PBE metadata.
   */
  static String encryptKey(String key, String password) throws KeyczarException {
    PBEKeyczarKey pbeMetadata = new PBEKeyczarKey();
    pbeMetadata.cipher = PBE_CIPHER;
    pbeMetadata.hmac = PBE_HMAC;
    pbeMetadata.iterationCount = DEFAULT_ITERATION_COUNT;
    // TODO: Figure out how to remove these base64 encoding round trips?
    byte[] salt = new byte[SALT_BYTES];
    Util.rand(salt);
    pbeMetadata.salt = Base64Coder.encodeWebSafe(salt);

    byte[] ivBytes = new byte[PBE_AES_KEY_BYTES];
    Util.rand(ivBytes);
    pbeMetadata.iv = Base64Coder.encodeWebSafe(ivBytes);

    try {
      Cipher encryptingCipher = initializePassphraseCipher(password, pbeMetadata, Cipher.ENCRYPT_MODE);
          Cipher.getInstance(CipherMode.CBC.getMode());

      byte[] keyBytes = key.getBytes("UTF-8");
      byte[] encrypted = encryptingCipher.doFinal(keyBytes);

      pbeMetadata.key = Base64Coder.encodeWebSafe(encrypted);
    } catch (java.security.GeneralSecurityException e) {
      throw new KeyczarException("Error encrypting key", e);
    } catch (UnsupportedEncodingException e) {
      throw new RuntimeException("Should never occur");
    }

    return Util.gson().toJson(pbeMetadata);
  }

  static PBEKeySpec keySpecFromJson(PBEKeyczarKey pbeKey, String password) {
    try {
      byte[] saltBytes = Base64Coder.decodeWebSafe(pbeKey.salt);
      char[] passwordChars = password.toCharArray();
      // keyLength is in bits
      return new PBEKeySpec(passwordChars, saltBytes, pbeKey.iterationCount, PBE_AES_KEY_BYTES*8);
    } catch (Base64DecodingException e) {
      throw new RuntimeException(e);
    }
  }

  private static final String PBE_CIPHER = "AES128";
  // TODO: Look this up from the cipher type
  static final int PBE_AES_KEY_BYTES = 16;
  private static final String PBE_HMAC = "HMAC_SHA1";

  static PBEKeyczarKey parsePBEMetadata(String json) throws KeyczarException {
    PBEKeyczarKey pbeKey = Util.gson().fromJson(json, PBEKeyczarKey.class);
    if (pbeKey.cipher == null || pbeKey.hmac == null || pbeKey.iterationCount == 0) {
      throw new KeyczarException("Invalid password-based key metadata");
    }

    if (!pbeKey.cipher.equals(PBE_CIPHER)) {
      throw new KeyczarException("Unsupported password-based encryption algorithm");
    }
    if (!pbeKey.hmac.equals(PBE_HMAC)) {
      throw new KeyczarException("Unsupported key derivation hash function");
    }
    if (pbeKey.iterationCount <= 0) {
      throw new KeyczarException("Invalid iteration count");
    }
    return pbeKey;
  }
}
